# 编写逆扩散采样函数，从xt恢复x0，其需要五个参数：model（模型），shape（形状）， n_steps（时间步数），betas和
# one_minus_alphas_bar_sqrt（这是一个张量，包含计算中使用的与（1-alpha）相关的值）
import torch


def p_sample_loop(model, shape, n_steps, betas, one_minus_alphas_bar_sqrt):
    """从x[T]恢复x[T-1]、x[T-2]|...x[0]"""
    # 使用torch.randn函数生成一个形状为shape的随机张量cur_x
    cur_x = torch.randn(shape)
    # 将cur_x添加到x_seq列表中
    x_seq = [cur_x]
    # 使用reversed(range(n_steps))生成一个倒序的迭代器，循环遍历每个时间步
    for i in reversed(range(n_steps)):
        # 在每个时间步上，调用p_sample函数，并将model、cur_x、i、betas和one_minus_alphas_bar_sqrt作为参数传递给它，p_sample函数用于采样恢复当前时间步的重构值
        cur_x = p_sample(model, cur_x, i, betas, one_minus_alphas_bar_sqrt)
        # 将重构值cur_x添加到x_seq列表中
        x_seq.append(cur_x)
    # 返回保存了所有时间步采样结果的列表x_seq作为函数的输出结果
    return x_seq


'''p_sample_loop函数实现了一个循环，通过反向遍历时间步，从最后一个时间步开始逐步恢复重构值，并将每个时间步的重构值添加到列表中。最后，返回保存了所有时间步重构值的列表。这个过程用于从最后一个时间步恢复到初始状态的逆扩散采样过程'''


# p_sample函数用于采样恢复当前时间步的重构值，其需要五个参数：model（模型），x（输入），t（时间步），betas和one_minus_alphas_bar_sqrt（这是一个张量，包含计算中使用的与（1-alpha）相关的值）
def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    """从x[T]采样t时刻的重构值"""
    # 将t转换为一个张量torch.tensor([t])
    t = torch.tensor([t])

    # 计算系数coeff，其值为betas[t]除以one_minus_alphas_bar_sqrt[t]，这些参数用于计算采样过程中的均值
    coeff = betas[t] / one_minus_alphas_bar_sqrt[t]

    # 调用model函数，将输入x和时间步t作为参数传递给它，返回一个张量eps_theta，该张量表示模型在给定输入和时间步的情况下产生的噪声
    eps_theta = model(x, t)

    # 计算均值mean，该均值受到输入x、系数betas、one_minus_alphas_bar_sqrt以及模型产生的噪声eps_theta的影响
    mean = (1 / (1 - betas[t]).sqrt()) * (x - (coeff * eps_theta))

    # 生成一个与输入x具有相同形状的随机张量z，使用torch.randn_like(x)来实现
    z = torch.randn_like(x)

    # 计算标准差sigma_t，其值为betas[t]的平方根
    sigma_t = betas[t].sqrt()

    # 计算最终的采样结果，得到sample的分布
    sample = mean + sigma_t * z

    # 返回sample的分布
    return (sample)

    '''p_sample函数实现了采样恢复当前时间步的重构值'''